{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([160, 250, 256])\n",
      "torch.Size([160, 250, 256])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "max_len = 5000\n",
    "base = 10000\n",
    "num_features = 256\n",
    "\n",
    "x = torch.zeros(250, 5 * 32, 256)\n",
    "position = torch.arange(max_len)  # (max_len,)\n",
    "index = torch.arange(0, num_features, 2) / num_features  # (num_features // 2,)\n",
    "indices = position.unsqueeze(dim=1) / (\n",
    "    base ** index.unsqueeze(dim=0)\n",
    ")  # (max_len, num_features // 2)\n",
    "sin, cos = torch.sin(indices), torch.cos(indices)\n",
    "positional_encoding = torch.stack([sin, cos], dim=-1)  # (max_len, num_features // 2, 2)\n",
    "pos = positional_encoding.view(max_len, 1, num_features)\n",
    "out = x + pos[:x.size(0)]\n",
    "out = out.permute(1,0,2)\n",
    "\n",
    "\n",
    "print(out.shape)\n",
    "\n",
    "import math\n",
    "x = torch.zeros(5 * 32, 250, 256)\n",
    "input_size = 256\n",
    "pe = torch.zeros(max_len, input_size, requires_grad=False)\n",
    "positions = torch.arange(0, max_len).unsqueeze(1).float()\n",
    "denominator = torch.exp(\n",
    "    torch.arange(0, input_size, 2).float() * -(math.log(10000.0) / input_size)\n",
    ")\n",
    "\n",
    "pe[:, 0::2] = torch.sin(positions * denominator)\n",
    "pe[:, 1::2] = torch.cos(positions * denominator)\n",
    "pe = pe\n",
    "\n",
    "pes = pe[:x.size(1),:]\n",
    "out2 = pes + x\n",
    "print(out2.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.0000e+00,  1.0000e+00,  0.0000e+00,  ...,  1.0000e+00,\n",
       "           0.0000e+00,  1.0000e+00],\n",
       "         [ 8.4147e-01,  5.4030e-01,  8.0196e-01,  ...,  1.0000e+00,\n",
       "           1.0746e-04,  1.0000e+00],\n",
       "         [ 9.0930e-01, -4.1615e-01,  9.5814e-01,  ...,  1.0000e+00,\n",
       "           2.1492e-04,  1.0000e+00],\n",
       "         ...,\n",
       "         [ 9.2681e-01, -3.7554e-01, -4.9256e-01,  ...,  9.9959e-01,\n",
       "           2.6540e-02,  9.9965e-01],\n",
       "         [ 1.8475e-01, -9.8279e-01, -9.9217e-01,  ...,  9.9959e-01,\n",
       "           2.6647e-02,  9.9964e-01],\n",
       "         [-7.2716e-01, -6.8646e-01, -6.9284e-01,  ...,  9.9959e-01,\n",
       "           2.6755e-02,  9.9964e-01]],\n",
       "\n",
       "        [[ 0.0000e+00,  1.0000e+00,  0.0000e+00,  ...,  1.0000e+00,\n",
       "           0.0000e+00,  1.0000e+00],\n",
       "         [ 8.4147e-01,  5.4030e-01,  8.0196e-01,  ...,  1.0000e+00,\n",
       "           1.0746e-04,  1.0000e+00],\n",
       "         [ 9.0930e-01, -4.1615e-01,  9.5814e-01,  ...,  1.0000e+00,\n",
       "           2.1492e-04,  1.0000e+00],\n",
       "         ...,\n",
       "         [ 9.2681e-01, -3.7554e-01, -4.9256e-01,  ...,  9.9959e-01,\n",
       "           2.6540e-02,  9.9965e-01],\n",
       "         [ 1.8475e-01, -9.8279e-01, -9.9217e-01,  ...,  9.9959e-01,\n",
       "           2.6647e-02,  9.9964e-01],\n",
       "         [-7.2716e-01, -6.8646e-01, -6.9284e-01,  ...,  9.9959e-01,\n",
       "           2.6755e-02,  9.9964e-01]],\n",
       "\n",
       "        [[ 0.0000e+00,  1.0000e+00,  0.0000e+00,  ...,  1.0000e+00,\n",
       "           0.0000e+00,  1.0000e+00],\n",
       "         [ 8.4147e-01,  5.4030e-01,  8.0196e-01,  ...,  1.0000e+00,\n",
       "           1.0746e-04,  1.0000e+00],\n",
       "         [ 9.0930e-01, -4.1615e-01,  9.5814e-01,  ...,  1.0000e+00,\n",
       "           2.1492e-04,  1.0000e+00],\n",
       "         ...,\n",
       "         [ 9.2681e-01, -3.7554e-01, -4.9256e-01,  ...,  9.9959e-01,\n",
       "           2.6540e-02,  9.9965e-01],\n",
       "         [ 1.8475e-01, -9.8279e-01, -9.9217e-01,  ...,  9.9959e-01,\n",
       "           2.6647e-02,  9.9964e-01],\n",
       "         [-7.2716e-01, -6.8646e-01, -6.9284e-01,  ...,  9.9959e-01,\n",
       "           2.6755e-02,  9.9964e-01]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[ 0.0000e+00,  1.0000e+00,  0.0000e+00,  ...,  1.0000e+00,\n",
       "           0.0000e+00,  1.0000e+00],\n",
       "         [ 8.4147e-01,  5.4030e-01,  8.0196e-01,  ...,  1.0000e+00,\n",
       "           1.0746e-04,  1.0000e+00],\n",
       "         [ 9.0930e-01, -4.1615e-01,  9.5814e-01,  ...,  1.0000e+00,\n",
       "           2.1492e-04,  1.0000e+00],\n",
       "         ...,\n",
       "         [ 9.2681e-01, -3.7554e-01, -4.9256e-01,  ...,  9.9959e-01,\n",
       "           2.6540e-02,  9.9965e-01],\n",
       "         [ 1.8475e-01, -9.8279e-01, -9.9217e-01,  ...,  9.9959e-01,\n",
       "           2.6647e-02,  9.9964e-01],\n",
       "         [-7.2716e-01, -6.8646e-01, -6.9284e-01,  ...,  9.9959e-01,\n",
       "           2.6755e-02,  9.9964e-01]],\n",
       "\n",
       "        [[ 0.0000e+00,  1.0000e+00,  0.0000e+00,  ...,  1.0000e+00,\n",
       "           0.0000e+00,  1.0000e+00],\n",
       "         [ 8.4147e-01,  5.4030e-01,  8.0196e-01,  ...,  1.0000e+00,\n",
       "           1.0746e-04,  1.0000e+00],\n",
       "         [ 9.0930e-01, -4.1615e-01,  9.5814e-01,  ...,  1.0000e+00,\n",
       "           2.1492e-04,  1.0000e+00],\n",
       "         ...,\n",
       "         [ 9.2681e-01, -3.7554e-01, -4.9256e-01,  ...,  9.9959e-01,\n",
       "           2.6540e-02,  9.9965e-01],\n",
       "         [ 1.8475e-01, -9.8279e-01, -9.9217e-01,  ...,  9.9959e-01,\n",
       "           2.6647e-02,  9.9964e-01],\n",
       "         [-7.2716e-01, -6.8646e-01, -6.9284e-01,  ...,  9.9959e-01,\n",
       "           2.6755e-02,  9.9964e-01]],\n",
       "\n",
       "        [[ 0.0000e+00,  1.0000e+00,  0.0000e+00,  ...,  1.0000e+00,\n",
       "           0.0000e+00,  1.0000e+00],\n",
       "         [ 8.4147e-01,  5.4030e-01,  8.0196e-01,  ...,  1.0000e+00,\n",
       "           1.0746e-04,  1.0000e+00],\n",
       "         [ 9.0930e-01, -4.1615e-01,  9.5814e-01,  ...,  1.0000e+00,\n",
       "           2.1492e-04,  1.0000e+00],\n",
       "         ...,\n",
       "         [ 9.2681e-01, -3.7554e-01, -4.9256e-01,  ...,  9.9959e-01,\n",
       "           2.6540e-02,  9.9965e-01],\n",
       "         [ 1.8475e-01, -9.8279e-01, -9.9217e-01,  ...,  9.9959e-01,\n",
       "           2.6647e-02,  9.9964e-01],\n",
       "         [-7.2716e-01, -6.8646e-01, -6.9284e-01,  ...,  9.9959e-01,\n",
       "           2.6755e-02,  9.9964e-01]]])"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.0000e+00,  1.0000e+00,  0.0000e+00,  ...,  1.0000e+00,\n",
       "           0.0000e+00,  1.0000e+00],\n",
       "         [ 8.4147e-01,  5.4030e-01,  8.0196e-01,  ...,  1.0000e+00,\n",
       "           1.0746e-04,  1.0000e+00],\n",
       "         [ 9.0930e-01, -4.1615e-01,  9.5814e-01,  ...,  1.0000e+00,\n",
       "           2.1492e-04,  1.0000e+00],\n",
       "         ...,\n",
       "         [ 9.2681e-01, -3.7554e-01, -4.9256e-01,  ...,  9.9959e-01,\n",
       "           2.6540e-02,  9.9965e-01],\n",
       "         [ 1.8475e-01, -9.8279e-01, -9.9217e-01,  ...,  9.9959e-01,\n",
       "           2.6647e-02,  9.9964e-01],\n",
       "         [-7.2716e-01, -6.8646e-01, -6.9284e-01,  ...,  9.9959e-01,\n",
       "           2.6755e-02,  9.9964e-01]],\n",
       "\n",
       "        [[ 0.0000e+00,  1.0000e+00,  0.0000e+00,  ...,  1.0000e+00,\n",
       "           0.0000e+00,  1.0000e+00],\n",
       "         [ 8.4147e-01,  5.4030e-01,  8.0196e-01,  ...,  1.0000e+00,\n",
       "           1.0746e-04,  1.0000e+00],\n",
       "         [ 9.0930e-01, -4.1615e-01,  9.5814e-01,  ...,  1.0000e+00,\n",
       "           2.1492e-04,  1.0000e+00],\n",
       "         ...,\n",
       "         [ 9.2681e-01, -3.7554e-01, -4.9256e-01,  ...,  9.9959e-01,\n",
       "           2.6540e-02,  9.9965e-01],\n",
       "         [ 1.8475e-01, -9.8279e-01, -9.9217e-01,  ...,  9.9959e-01,\n",
       "           2.6647e-02,  9.9964e-01],\n",
       "         [-7.2716e-01, -6.8646e-01, -6.9284e-01,  ...,  9.9959e-01,\n",
       "           2.6755e-02,  9.9964e-01]],\n",
       "\n",
       "        [[ 0.0000e+00,  1.0000e+00,  0.0000e+00,  ...,  1.0000e+00,\n",
       "           0.0000e+00,  1.0000e+00],\n",
       "         [ 8.4147e-01,  5.4030e-01,  8.0196e-01,  ...,  1.0000e+00,\n",
       "           1.0746e-04,  1.0000e+00],\n",
       "         [ 9.0930e-01, -4.1615e-01,  9.5814e-01,  ...,  1.0000e+00,\n",
       "           2.1492e-04,  1.0000e+00],\n",
       "         ...,\n",
       "         [ 9.2681e-01, -3.7554e-01, -4.9256e-01,  ...,  9.9959e-01,\n",
       "           2.6540e-02,  9.9965e-01],\n",
       "         [ 1.8475e-01, -9.8279e-01, -9.9217e-01,  ...,  9.9959e-01,\n",
       "           2.6647e-02,  9.9964e-01],\n",
       "         [-7.2716e-01, -6.8646e-01, -6.9284e-01,  ...,  9.9959e-01,\n",
       "           2.6755e-02,  9.9964e-01]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[ 0.0000e+00,  1.0000e+00,  0.0000e+00,  ...,  1.0000e+00,\n",
       "           0.0000e+00,  1.0000e+00],\n",
       "         [ 8.4147e-01,  5.4030e-01,  8.0196e-01,  ...,  1.0000e+00,\n",
       "           1.0746e-04,  1.0000e+00],\n",
       "         [ 9.0930e-01, -4.1615e-01,  9.5814e-01,  ...,  1.0000e+00,\n",
       "           2.1492e-04,  1.0000e+00],\n",
       "         ...,\n",
       "         [ 9.2681e-01, -3.7554e-01, -4.9256e-01,  ...,  9.9959e-01,\n",
       "           2.6540e-02,  9.9965e-01],\n",
       "         [ 1.8475e-01, -9.8279e-01, -9.9217e-01,  ...,  9.9959e-01,\n",
       "           2.6647e-02,  9.9964e-01],\n",
       "         [-7.2716e-01, -6.8646e-01, -6.9284e-01,  ...,  9.9959e-01,\n",
       "           2.6755e-02,  9.9964e-01]],\n",
       "\n",
       "        [[ 0.0000e+00,  1.0000e+00,  0.0000e+00,  ...,  1.0000e+00,\n",
       "           0.0000e+00,  1.0000e+00],\n",
       "         [ 8.4147e-01,  5.4030e-01,  8.0196e-01,  ...,  1.0000e+00,\n",
       "           1.0746e-04,  1.0000e+00],\n",
       "         [ 9.0930e-01, -4.1615e-01,  9.5814e-01,  ...,  1.0000e+00,\n",
       "           2.1492e-04,  1.0000e+00],\n",
       "         ...,\n",
       "         [ 9.2681e-01, -3.7554e-01, -4.9256e-01,  ...,  9.9959e-01,\n",
       "           2.6540e-02,  9.9965e-01],\n",
       "         [ 1.8475e-01, -9.8279e-01, -9.9217e-01,  ...,  9.9959e-01,\n",
       "           2.6647e-02,  9.9964e-01],\n",
       "         [-7.2716e-01, -6.8646e-01, -6.9284e-01,  ...,  9.9959e-01,\n",
       "           2.6755e-02,  9.9964e-01]],\n",
       "\n",
       "        [[ 0.0000e+00,  1.0000e+00,  0.0000e+00,  ...,  1.0000e+00,\n",
       "           0.0000e+00,  1.0000e+00],\n",
       "         [ 8.4147e-01,  5.4030e-01,  8.0196e-01,  ...,  1.0000e+00,\n",
       "           1.0746e-04,  1.0000e+00],\n",
       "         [ 9.0930e-01, -4.1615e-01,  9.5814e-01,  ...,  1.0000e+00,\n",
       "           2.1492e-04,  1.0000e+00],\n",
       "         ...,\n",
       "         [ 9.2681e-01, -3.7554e-01, -4.9256e-01,  ...,  9.9959e-01,\n",
       "           2.6540e-02,  9.9965e-01],\n",
       "         [ 1.8475e-01, -9.8279e-01, -9.9217e-01,  ...,  9.9959e-01,\n",
       "           2.6647e-02,  9.9964e-01],\n",
       "         [-7.2716e-01, -6.8646e-01, -6.9284e-01,  ...,  9.9959e-01,\n",
       "           2.6755e-02,  9.9964e-01]]])"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 1.4693e+00,  2.2691e-01, -9.1297e-01,  ...,  2.2378e+00,\n",
       "           3.2382e-01, -7.2300e-01],\n",
       "         [-7.2954e-01, -6.0545e-01,  2.3867e+00,  ...,  1.3786e+00,\n",
       "          -1.8692e+00,  1.9029e+00],\n",
       "         [ 2.6610e+00, -8.8740e-01,  1.3546e+00,  ...,  1.7268e+00,\n",
       "           1.0024e+00, -5.2680e-01],\n",
       "         ...,\n",
       "         [ 2.8914e+00, -4.9321e-01, -1.4658e+00,  ...,  6.0006e-01,\n",
       "           1.8953e+00,  2.9899e-01],\n",
       "         [ 7.8183e-01, -9.8246e-01, -1.5234e+00,  ..., -1.8092e-01,\n",
       "           6.2749e-01,  1.7362e+00],\n",
       "         [-9.5553e-01, -6.0365e-01,  1.9070e-02,  ...,  2.5541e+00,\n",
       "          -8.2313e-01,  1.7813e-01]],\n",
       "\n",
       "        [[ 1.3095e+00,  1.2275e+00,  5.3896e-01,  ..., -9.5510e-01,\n",
       "           1.1240e-01,  2.0881e+00],\n",
       "         [ 1.6767e+00, -2.6069e-01,  9.2963e-01,  ...,  1.3396e+00,\n",
       "          -1.8916e+00,  3.5407e+00],\n",
       "         [-1.1947e-01, -4.4152e-01,  4.5651e-01,  ...,  6.3425e-02,\n",
       "           4.5486e-01,  8.7600e-01],\n",
       "         ...,\n",
       "         [ 7.5446e-01, -1.6513e-01, -2.4966e+00,  ...,  1.4789e+00,\n",
       "          -7.3630e-01,  2.2674e+00],\n",
       "         [ 4.4655e-01, -1.8776e+00,  2.8078e-01,  ...,  6.5701e-01,\n",
       "          -7.1061e-01,  1.4339e+00],\n",
       "         [-1.7774e+00, -2.5731e-01, -1.7468e+00,  ...,  2.4774e-01,\n",
       "          -2.8880e-01,  2.3193e+00]],\n",
       "\n",
       "        [[-1.4486e+00,  2.4510e+00, -1.5268e-03,  ...,  1.1448e+00,\n",
       "           1.2428e+00,  2.6312e+00],\n",
       "         [ 2.6484e+00,  1.9009e+00,  7.0590e-01,  ...,  1.0974e+00,\n",
       "          -3.1768e+00,  1.4168e+00],\n",
       "         [ 1.7511e+00, -8.5598e-01,  2.2435e+00,  ...,  6.6552e-01,\n",
       "           7.7726e-01,  1.6315e+00],\n",
       "         ...,\n",
       "         [ 2.6785e+00, -2.0185e+00, -6.9232e-01,  ...,  1.9242e+00,\n",
       "           1.1564e+00,  1.2052e+00],\n",
       "         [ 1.6294e+00, -2.4726e+00, -1.1752e+00,  ...,  6.1532e-01,\n",
       "           2.1202e+00, -7.5273e-02],\n",
       "         [-1.7392e+00, -7.4795e-02, -7.2993e-01,  ...,  7.1781e-01,\n",
       "          -7.5034e-01,  2.0936e+00]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[-3.0643e-01,  5.4783e-01,  7.0935e-01,  ...,  3.4184e+00,\n",
       "          -9.6318e-02, -8.9007e-01],\n",
       "         [-6.0018e-01, -1.0016e+00, -2.0606e-02,  ...,  4.6536e-01,\n",
       "           2.7667e-01,  3.7066e-01],\n",
       "         [ 1.9014e-01,  9.1118e-01, -2.9936e-01,  ...,  2.1033e+00,\n",
       "          -2.7853e-01,  1.8028e+00],\n",
       "         ...,\n",
       "         [ 9.3297e-01,  2.9057e-01, -6.7830e-05,  ...,  2.1142e-01,\n",
       "           2.0406e+00, -2.1218e-01],\n",
       "         [ 6.1064e-02,  6.3324e-01, -2.5048e+00,  ...,  4.7794e-01,\n",
       "          -8.7283e-01, -3.7227e-01],\n",
       "         [-1.2694e+00, -1.1043e+00,  9.6519e-01,  ...,  5.2248e-02,\n",
       "          -1.4083e-01,  1.0069e+00]],\n",
       "\n",
       "        [[ 5.3183e-01,  2.2185e+00,  1.6388e+00,  ...,  1.1581e+00,\n",
       "          -1.7711e-01,  1.1111e+00],\n",
       "         [ 1.8946e+00,  6.5683e-01, -1.0219e+00,  ...,  2.3801e+00,\n",
       "          -1.0939e-01,  1.3749e+00],\n",
       "         [ 1.6948e-01, -4.3153e-01,  1.1498e+00,  ...,  4.7027e-01,\n",
       "          -3.0196e-01,  1.6171e+00],\n",
       "         ...,\n",
       "         [-1.8807e-01, -4.3345e-01,  1.7584e-01,  ..., -8.8140e-01,\n",
       "          -2.5494e-01,  2.4830e+00],\n",
       "         [ 2.1569e-01,  3.0638e-02, -2.4070e+00,  ...,  3.0138e-01,\n",
       "          -1.1767e-01,  1.0513e+00],\n",
       "         [-1.0292e+00, -3.4056e+00, -3.0239e-01,  ...,  1.0230e+00,\n",
       "           1.1419e+00, -2.3560e-02]],\n",
       "\n",
       "        [[ 3.4135e-01,  1.2445e+00, -1.3996e+00,  ...,  8.2658e-01,\n",
       "          -1.2203e+00,  9.4308e-01],\n",
       "         [ 3.0357e-02,  2.6506e-01, -2.6207e+00,  ...,  1.2082e+00,\n",
       "           2.0959e+00,  7.5211e-01],\n",
       "         [ 1.2162e+00,  6.4184e-01,  1.4517e+00,  ...,  5.3494e-01,\n",
       "          -1.4623e+00,  7.2473e-01],\n",
       "         ...,\n",
       "         [-1.2850e+00, -1.3199e+00,  1.0889e-01,  ...,  4.9053e-01,\n",
       "           7.7231e-01,  1.9247e-01],\n",
       "         [-1.2609e+00, -1.1776e+00, -9.2954e-01,  ...,  1.1684e+00,\n",
       "          -5.7123e-01,  4.4560e-01],\n",
       "         [ 1.2205e-01, -5.3322e-01,  5.0334e-01,  ...,  1.7063e+00,\n",
       "           4.1593e-02,  2.8199e+00]]])"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pe[:,:x.size(1)]+ x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 1.4693e+00,  2.2691e-01, -9.1297e-01,  ...,  2.2378e+00,\n",
       "           3.2382e-01, -7.2300e-01],\n",
       "         [-7.2954e-01, -6.0545e-01,  2.3867e+00,  ...,  1.3786e+00,\n",
       "          -1.8692e+00,  1.9029e+00],\n",
       "         [ 2.6610e+00, -8.8740e-01,  1.3546e+00,  ...,  1.7268e+00,\n",
       "           1.0024e+00, -5.2680e-01],\n",
       "         ...,\n",
       "         [ 2.8914e+00, -4.9321e-01, -1.4658e+00,  ...,  6.0006e-01,\n",
       "           1.8953e+00,  2.9899e-01],\n",
       "         [ 7.8183e-01, -9.8246e-01, -1.5234e+00,  ..., -1.8092e-01,\n",
       "           6.2749e-01,  1.7362e+00],\n",
       "         [-9.5553e-01, -6.0365e-01,  1.9070e-02,  ...,  2.5541e+00,\n",
       "          -8.2313e-01,  1.7813e-01]],\n",
       "\n",
       "        [[ 1.3095e+00,  1.2275e+00,  5.3896e-01,  ..., -9.5510e-01,\n",
       "           1.1240e-01,  2.0881e+00],\n",
       "         [ 1.6767e+00, -2.6069e-01,  9.2963e-01,  ...,  1.3396e+00,\n",
       "          -1.8916e+00,  3.5407e+00],\n",
       "         [-1.1947e-01, -4.4152e-01,  4.5651e-01,  ...,  6.3425e-02,\n",
       "           4.5486e-01,  8.7600e-01],\n",
       "         ...,\n",
       "         [ 7.5446e-01, -1.6513e-01, -2.4966e+00,  ...,  1.4789e+00,\n",
       "          -7.3630e-01,  2.2674e+00],\n",
       "         [ 4.4655e-01, -1.8776e+00,  2.8078e-01,  ...,  6.5701e-01,\n",
       "          -7.1061e-01,  1.4339e+00],\n",
       "         [-1.7774e+00, -2.5731e-01, -1.7468e+00,  ...,  2.4774e-01,\n",
       "          -2.8880e-01,  2.3193e+00]],\n",
       "\n",
       "        [[-1.4486e+00,  2.4510e+00, -1.5268e-03,  ...,  1.1448e+00,\n",
       "           1.2428e+00,  2.6312e+00],\n",
       "         [ 2.6484e+00,  1.9009e+00,  7.0590e-01,  ...,  1.0974e+00,\n",
       "          -3.1768e+00,  1.4168e+00],\n",
       "         [ 1.7511e+00, -8.5598e-01,  2.2435e+00,  ...,  6.6552e-01,\n",
       "           7.7726e-01,  1.6315e+00],\n",
       "         ...,\n",
       "         [ 2.6785e+00, -2.0185e+00, -6.9232e-01,  ...,  1.9242e+00,\n",
       "           1.1564e+00,  1.2052e+00],\n",
       "         [ 1.6294e+00, -2.4726e+00, -1.1752e+00,  ...,  6.1532e-01,\n",
       "           2.1202e+00, -7.5273e-02],\n",
       "         [-1.7392e+00, -7.4795e-02, -7.2993e-01,  ...,  7.1781e-01,\n",
       "          -7.5034e-01,  2.0936e+00]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[-3.0643e-01,  5.4783e-01,  7.0935e-01,  ...,  3.4184e+00,\n",
       "          -9.6318e-02, -8.9007e-01],\n",
       "         [-6.0018e-01, -1.0016e+00, -2.0606e-02,  ...,  4.6536e-01,\n",
       "           2.7667e-01,  3.7066e-01],\n",
       "         [ 1.9014e-01,  9.1118e-01, -2.9936e-01,  ...,  2.1033e+00,\n",
       "          -2.7853e-01,  1.8028e+00],\n",
       "         ...,\n",
       "         [ 9.3297e-01,  2.9057e-01, -6.7830e-05,  ...,  2.1142e-01,\n",
       "           2.0406e+00, -2.1218e-01],\n",
       "         [ 6.1064e-02,  6.3324e-01, -2.5048e+00,  ...,  4.7794e-01,\n",
       "          -8.7283e-01, -3.7227e-01],\n",
       "         [-1.2694e+00, -1.1043e+00,  9.6519e-01,  ...,  5.2248e-02,\n",
       "          -1.4083e-01,  1.0069e+00]],\n",
       "\n",
       "        [[ 5.3183e-01,  2.2185e+00,  1.6388e+00,  ...,  1.1581e+00,\n",
       "          -1.7711e-01,  1.1111e+00],\n",
       "         [ 1.8946e+00,  6.5683e-01, -1.0219e+00,  ...,  2.3801e+00,\n",
       "          -1.0939e-01,  1.3749e+00],\n",
       "         [ 1.6948e-01, -4.3153e-01,  1.1498e+00,  ...,  4.7027e-01,\n",
       "          -3.0196e-01,  1.6171e+00],\n",
       "         ...,\n",
       "         [-1.8807e-01, -4.3345e-01,  1.7584e-01,  ..., -8.8140e-01,\n",
       "          -2.5494e-01,  2.4830e+00],\n",
       "         [ 2.1569e-01,  3.0638e-02, -2.4070e+00,  ...,  3.0138e-01,\n",
       "          -1.1767e-01,  1.0513e+00],\n",
       "         [-1.0292e+00, -3.4056e+00, -3.0239e-01,  ...,  1.0230e+00,\n",
       "           1.1419e+00, -2.3560e-02]],\n",
       "\n",
       "        [[ 3.4135e-01,  1.2445e+00, -1.3996e+00,  ...,  8.2658e-01,\n",
       "          -1.2203e+00,  9.4308e-01],\n",
       "         [ 3.0357e-02,  2.6506e-01, -2.6207e+00,  ...,  1.2082e+00,\n",
       "           2.0959e+00,  7.5211e-01],\n",
       "         [ 1.2162e+00,  6.4184e-01,  1.4517e+00,  ...,  5.3494e-01,\n",
       "          -1.4623e+00,  7.2473e-01],\n",
       "         ...,\n",
       "         [-1.2850e+00, -1.3199e+00,  1.0889e-01,  ...,  4.9053e-01,\n",
       "           7.7231e-01,  1.9247e-01],\n",
       "         [-1.2609e+00, -1.1776e+00, -9.2954e-01,  ...,  1.1684e+00,\n",
       "          -5.7123e-01,  4.4560e-01],\n",
       "         [ 1.2205e-01, -5.3322e-01,  5.0334e-01,  ...,  1.7063e+00,\n",
       "           4.1593e-02,  2.8199e+00]]])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
